import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import os


def visualize_training(model_name):
    # 加载日志
    log_path = f'models/{model_name}/log.csv'
    log = pd.read_csv(log_path)

    # 设置绘图风格
    sns.set_style("whitegrid")
    plt.figure(figsize=(15, 10))

    # 绘制损失曲线
    plt.subplot(2, 2, 1)
    plt.plot(log['epoch'], log['loss'], label='Train Loss')
    plt.plot(log['epoch'], log['val_loss'], label='Validation Loss')
    plt.title('Loss Curve')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()

    # 绘制IoU曲线
    plt.subplot(2, 2, 2)
    plt.plot(log['epoch'], log['iou'], label='Train IoU')
    plt.plot(log['epoch'], log['val_iou'], label='Validation IoU')
    plt.title('IoU Curve')
    plt.xlabel('Epoch')
    plt.ylabel('IoU')
    plt.legend()

    # 绘制Dice曲线
    plt.subplot(2, 2, 3)
    plt.plot(log['epoch'], log['dice'], label='Train Dice')
    plt.plot(log['epoch'], log['val_dice'], label='Validation Dice')
    plt.title('Dice Coefficient Curve')
    plt.xlabel('Epoch')
    plt.ylabel('Dice')
    plt.legend()

    # 绘制准确率曲线
    plt.subplot(2, 2, 4)
    plt.plot(log['epoch'], log['acc'], label='Train Accuracy')
    plt.plot(log['epoch'], log['val_acc'], label='Validation Accuracy')
    plt.title('Accuracy Curve')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()

    plt.tight_layout()
    plt.savefig(f'models/{model_name}/training_curves.png', dpi=300)
    plt.show()


if __name__ == '__main__':
    model_name = 'dice_training'  # 替换为您的模型名称
    visualize_training(model_name)